# coding : utf-8

from scipy.spatial import KDTree

from collections import deque
import math

import carla
from agents.tools.misc import compute_distance

def get_transform_by_bias(base_transform, bias):
    """
    计算沿着垂直于给定transform点的yaw角度，偏移bias后，形成的新transform并返回
    坐标系为右前上，右侧bias为正，左侧bias为负。
    @param transform : carla.Transform
    @param bias : float
    @return new_transform : carla.Transform
    """
    new_tf = carla.Transform(rotation = base_transform.rotation,\
                             location = base_transform.location)
    # calc x,y
    theta = math.radians(base_transform.rotation.yaw)
    new_tf.location.x = base_transform.location.x - bias * math.sin(theta) # TODO
    new_tf.location.y = base_transform.location.y + bias * math.cos(theta) # TODO
    return new_tf

class FrenetPoint:
    """
    .centre_location 中心点位置
    .tansform 其中的location为本点所在的位置，而不是中心点位置
    .s 
    .l 中心偏移量
    .kappa 曲率
    .dappa 曲率变化率
    """
    def __init__(self, centre_tf:carla.Transform, centre_s=0.0, bias=0.0, kappa=0.0, dappa=0.0):
        """
        @param transform : 车道中心线的transform
        @param centre_s : 车道中心线的s
        @param bias : 相对车道中心线的偏移，符合右手坐标系，右前上规则
        @param kappa : 车道中心线的曲率
        @param dappa : 车道中心线的曲率变化率
        """
        self.centre_location = centre_tf.location
        self.transform = get_transform_by_bias(centre_tf, bias)
        self.l = bias 
        if kappa == 0:
            self.s = centre_s # 同比例缩小/放大
            self.kappa = kappa
            self.dappa = dappa
        else:
            radius0 = 1 / kappa
            radius1 = radius0 + bias
            self.s = centre_s * radius1 / radius0 # 同比例缩小/放大
            self.kappa = 1 / radius1
            self.dappa = dappa # 朝向不变，dappa就不变

    # def __init__(self, ref_frenetpoint, bias):
    #     self.__init__(ref_frenetpoint.transform, ref_frenetpoint.s, bias, ref_frenetpoint.kappa, ref_frenetpoint.dappa)
    
    def get_certre_tf(self):
        centre_tf = carla.Transform()
        centre_tf.rotation = self.transfom.rotation
        centre_tf.location = self.centre_location
        return centre_tf

    def get_nearest_waypoint(self, g_map:carla.Map, project_to_road = False):
        """
        @param g_map : carla.Map
        @param project_to_road : bool default = false
        @return waypoint : carla.Waypoint
        """
        return g_map.get_waypoint(self.transform.location, project_to_road)

class LocalRoadTreeDeque:
    """
    利用kd tree存储局部路径transform的location数据
    因为如果直接用transform数据来建立kdtree的话，会多余rpy等角度数据
    实际使用中是单点对单点的寻找，所以利用本类来构造只包含location的tree
    但同时返回transform数据
    总体为deque,包含了frenetpoint_deque 和 location_deque
    """
    def __init__(self, maxlen = 1):
        # super(LocalRoadTreeDeque, self).__init__(maxlen=maxlen)
        self._local_transform_buffer = deque(maxlen=maxlen)
        self._local_location_buffer = deque(maxlen=maxlen)

    def set_tree(self):
        self._tree = KDTree(self._local_location_buffer)

    def popleft(self):
        self._local_transform_buffer.popleft()
        self._local_location_buffer.popleft()

    def append(self, local_data:FrenetPoint):
        """
        @param local_data : FrenetPoint 
        """
        self._local_transform_buffer.append(local_data)
        self._local_location_buffer.append((local_data.transform.location.x, 
            local_data.transform.location.y, 
            local_data.transform.location.z))

    def query(self, locations, k=1):
        """
        @param locations: carla.Location or list of carla.Location
        @parma k: the k nearest points will be queried, default is 1
        @return : dists , indexes
        """
        if isinstance(locations, list):
            query_data = []
            for loc in locations:
                query_data.append((loc.x, loc.y, loc.z))
        else:
            query_data = (locations.x, 
                locations.y, 
                locations.z)
        return self._tree.query(query_data, k=k)

    def clear(self):
        self._local_transform_buffer.clear()
        self._local_location_buffer.clear()

    def __len__(self):
        return len(self._local_transform_buffer)

    def __getitem__(self, index):
        return self._local_transform_buffer[index]
    
    def __iter__(self):
        return self._local_transform_buffer.__iter__()

    def __init_subclass__(self, index):
        return self._local_transform_buffer[index]
